import sys
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt


from colorama import Fore, Style
from numpy.lib.shape_base import expand_dims


def green(s: str):
    return f"{Fore.GREEN}{s}{Style.RESET_ALL}"

class Calculator:
    def __init__(self) -> None:
        self.pair = {}

    def build(self):
        print(green("请输入多项式："))
        print(green("例如，输入多项式 2*x^2 - x, 则输入（i-th系数、i-th指数）：2, 2, -1, 1"))
        s = input(green("> "))

        a = s.split(",")
        for factor, exp in zip(a[::2], a[1::2]):
            self.pair[eval(exp)] = self.pair.get(eval(exp), 0) + eval(factor)
        print(green("输入多项式：")+self.__repr__())
    
    def print(self, verbose=False):
        s = "f(x) = "
        l = len(s)
        # s = ""
        for exp, factor in sorted(self.pair.items(), key=lambda x:x[0], reverse=True):
            if factor < 0:
                pass
            else:
                if len(s) == l:
                    pass
                else:
                    s += "+ "
            s += f"{factor}*x^{exp}"
        # s = "f(x) = " + s[]
        if verbose:
            print(green(s))
        return s
    
    def get_x(self, x):
        a = 0
        for exp, factor in self.pair.items():
            try:
                a += factor * x ** exp
            except ZeroDivisionError:
                pass
        return a

    def plot(self, start=0, stop=1, step=0.01):
        xs = np.arange(start, stop, step)
        ys = np.apply_along_axis(arr=xs, func1d=self.get_x, axis=0)
        mpl.style.use("seaborn")
        l = self.print(verbose=False)
        if "^-" not in l:
            plt.plot(xs, ys, label=f"${l}$".replace("*", ""))
        else:
            plt.plot(xs, ys, label=f"{l}".replace("*", ""))
        plt.legend()
        plt.show()
    
    def __repr__(self) -> str:
        return self.print()
    
    def __add__(self, other):
        c = Calculator()
        c.pair = {}
        for exp, factor in self.pair.items():
            c.pair[exp] = factor
        for exp, factor in other.pair.items():
            c.pair[exp] = c.pair.get(exp, 0) + factor
        return c

    def __sub__(self, other):
        c = Calculator()
        c.pair = {}
        for exp, factor in self.pair.items():
            c.pair[exp] = factor
        for exp, factor in other.pair.items():
            c.pair[exp] = c.pair.get(exp, 0) - factor
        return c

def c_print():
    c1 = Calculator()
    c1.build()
    print(c1)

def c_add():
    c1 = Calculator()
    c1.build()
    c2 = Calculator()
    c2.build()
    a = c1 + c2
    print(a)

def c_subs():
    c1 = Calculator()
    c1.build()
    c2 = Calculator()
    c2.build()
    a = c1 - c2
    print(a)

def c_get(x):
    c1 = Calculator()
    c1.build()
    print(c1.get_x(x=x))

def c_plot(start, end, step):
    c1 = Calculator()
    c1.build()
    c1.plot(start=start, stop=end, step=step)


if __name__ == "__main__":
    # c_print()
    # c_add()
    # c_subs()
    # c_get(1)
    c_plot(0.001, 5, 0.01)